# ============================================================================
# COMPLETE DetectionModel class with ALL code + StatA integration
# ============================================================================

class DetectionModel(BaseModel):
    # YOLO detection model
    def __init__(self, cfg='yolo.yaml', ch=3, nc=None, anchors=None):  # model, input channels, number of classes
        super().__init__()
        
        # ============================================================
        # EXISTING YOLO CODE: Parse configuration
        # ============================================================
        if isinstance(cfg, dict):
            self.yaml = cfg  # model dict
        else:  # is *.yaml
            import yaml  # for torch hub
            self.yaml_file = Path(cfg).name
            with open(cfg, encoding='ascii', errors='ignore') as f:
                self.yaml = yaml.safe_load(f)  # model dict

        # ============================================================
        # EXISTING YOLO CODE: Define model
        # ============================================================
        ch = self.yaml['ch'] = self.yaml.get('ch', ch)  # input channels
        if nc and nc != self.yaml['nc']:
            LOGGER.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}")
            self.yaml['nc'] = nc  # override yaml value
        if anchors:
            LOGGER.info(f'Overriding model.yaml anchors with anchors={anchors}')
            self.yaml['anchors'] = round(anchors)  # override yaml value
        self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch])  # model, savelist
        self.names = [str(i) for i in range(self.yaml['nc'])]  # default names
        self.inplace = self.yaml.get('inplace', True)

        # ============================================================
        # EXISTING YOLO CODE: Build strides, anchors
        # ============================================================
        m = self.model[-1]  # Detect()
        if isinstance(m, (Detect, DDetect, Segment, DSegment, Panoptic)):
            s = 256  # 2x min stride
            m.inplace = self.inplace
            forward = lambda x: self.forward(x)[0] if isinstance(m, (Segment, DSegment, Panoptic)) else self.forward(x)
            m.stride = torch.tensor([s / x.shape[-2] for x in forward(torch.zeros(1, ch, s, s))])  # forward
            # check_anchor_order(m)
            # m.anchors /= m.stride.view(-1, 1, 1)
            self.stride = m.stride
            m.bias_init()  # only run once
            
        if isinstance(m, (DualDetect, TripleDetect, DualDDetect, TripleDDetect, DualDSegment)):
            s = 256  # 2x min stride
            m.inplace = self.inplace
            forward = lambda x: self.forward(x)[0][0] if isinstance(m, (DualDSegment)) else self.forward(x)[0]
            m.stride = torch.tensor([s / x.shape[-2] for x in forward(torch.zeros(1, ch, s, s))])  # forward
            # check_anchor_order(m)
            # m.anchors /= m.stride.view(-1, 1, 1)
            self.stride = m.stride
            m.bias_init()  # only run once

        # ============================================================
        # EXISTING YOLO CODE: Init weights, biases
        # ============================================================
        initialize_weights(self)
        self.info()
        LOGGER.info('')
        
        # ============================================================
        # NEW: StatA integration attributes
        # ============================================================
        self.enable_stata = False
        self.clip_projector = None
        self.clip_prototypes = None

    # ============================================================
    # EXISTING YOLO METHODS (unchanged)
    # ============================================================
    
    def forward(self, x, augment=False, profile=False, visualize=False):
        if augment:
            return self._forward_augment(x)  # augmented inference, None
        return self._forward_once(x, profile, visualize)  # single-scale inference, train

    def _forward_augment(self, x):
        img_size = x.shape[-2:]  # height, width
        s = [1, 0.83, 0.67]  # scales
        f = [None, 3, None]  # flips (2-ud, 3-lr)
        y = []  # outputs
        for si, fi in zip(s, f):
            xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max()))
            yi = self._forward_once(xi)[0]  # forward
            # cv2.imwrite(f'img_{si}.jpg', 255 * xi[0].cpu().numpy().transpose((1, 2, 0))[:, :, ::-1])  # save
            yi = self._descale_pred(yi, fi, si, img_size)
            y.append(yi)
        y = self._clip_augmented(y)  # clip augmented tails
        return torch.cat(y, 1), None  # augmented inference, train

    def _descale_pred(self, p, flips, scale, img_size):
        # de-scale predictions following augmented inference (inverse operation)
        if self.inplace:
            p[..., :4] /= scale  # de-scale
            if flips == 2:
                p[..., 1] = img_size[0] - p[..., 1]  # de-flip ud
            elif flips == 3:
                p[..., 0] = img_size[1] - p[..., 0]  # de-flip lr
        else:
            x, y, wh = p[..., 0:1] / scale, p[..., 1:2] / scale, p[..., 2:4] / scale  # de-scale
            if flips == 2:
                y = img_size[0] - y  # de-flip ud
            elif flips == 3:
                x = img_size[1] - x  # de-flip lr
            p = torch.cat((x, y, wh, p[..., 4:]), -1)
        return p

    def _clip_augmented(self, y):
        # Clip YOLO augmented inference tails
        nl = self.model[-1].nl  # number of detection layers (P3-P5)
        g = sum(4 ** x for x in range(nl))  # grid points
        e = 1  # exclude layer count
        i = (y[0].shape[1] // g) * sum(4 ** x for x in range(e))  # indices
        y[0] = y[0][:, :-i]  # large
        i = (y[-1].shape[1] // g) * sum(4 ** (nl - 1 - x) for x in range(e))  # indices
        y[-1] = y[-1][:, i:]  # small
        return y
    
    # ============================================================
    # NEW: StatA methods
    # ============================================================
    
    def setup_stata(self, class_names, clip_model_name='ViT-B/16'):
        """
        Initialize StatA for test-time adaptation
        
        Args:
            class_names: List of class names (e.g., COCO classes)
            clip_model_name: CLIP model architecture ('ViT-B/16' recommended)
        
        Example:
            model = DetectionModel('yolov9-c.yaml')
            model.setup_stata(['person', 'car', 'dog'])
        """
        from utils.clip_integration import YOLOCLIPProjector, CLIPTextEncoder
        
        print("\n" + "="*60)
        print("Setting up StatA Test-Time Adaptation")
        print("="*60)
        
        # Enable StatA mode
        self.enable_stata = True
        device = next(self.parameters()).device
        
        # Create CLIP feature projector
        if self.clip_projector is None:
            # YOLOv9-C channel configuration: [P3, P4, P5]
            # For YOLOv9-E use: [384, 768, 1536]
            yolo_channels = [256, 512, 1024]
            
            print(f"[StatA] Creating projectors: {yolo_channels} → 512")
            self.clip_projector = YOLOCLIPProjector(
                yolo_channels=yolo_channels, 
                clip_dim=512
            ).to(device)
        
        # Generate CLIP text prototypes
        print(f"[StatA] Encoding {len(class_names)} classes...")
        clip_encoder = CLIPTextEncoder(clip_model_name, device)
        self.clip_prototypes = clip_encoder.encode_classes(class_names)
        
        print(f"[StatA] ✓ Ready with {len(class_names)} classes")
        print(f"[StatA] Text prototypes: {self.clip_prototypes.shape}")
        print("="*60 + "\n")
    
    def extract_features_for_stata(self, x):
        """
        Extract multi-scale features (P3, P4, P5) for StatA
        
        Args:
            x: Input tensor [B, 3, H, W]
        
        Returns:
            features: List of 3 tensors [P3, P4, P5]
                P3: [B, 256, 80, 80]   (fine details)
                P4: [B, 512, 40, 40]   (medium scale)
                P5: [B, 1024, 20, 20]  (coarse scale)
        
        Note:
            The layer indices (15, 18, 21) are for YOLOv9-C.
            If you get "Expected 3 scales, got X" warning, you need to find
            the correct indices for your model using the find_layers.py script.
        """
        features = []
        y = []
        
        # Forward through all layers
        for i, m in enumerate(self.model):
            # Handle layer connections
            if m.f != -1:  # if not from previous layer
                x = y[m.f] if isinstance(m.f, int) else \
                    [x if j == -1 else y[j] for j in m.f]
            
            # Forward through layer
            x = m(x)
            
            # Store output if needed for later layers
            y.append(x if m.i in self.save else None)
            
            # Collect P3, P4, P5 features
            # These indices work for YOLOv9-C
            # Adjust if your model uses different architecture
            if m.i == 15:  # P3 output layer
                features.append(x)
            elif m.i == 18:  # P4 output layer
                features.append(x)
            elif m.i == 21:  # P5 output layer
                features.append(x)
        
        # Verify we got all 3 scales
        if len(features) != 3:
            print(f"[WARNING] Expected 3 scales, got {len(features)}")
            print(f"[INFO] You may need to adjust layer indices in extract_features_for_stata()")
            print(f"[INFO] Run find_layers.py to find correct indices for your model")
        
        return features


# ============================================================================
# Keep this line unchanged at the end of the file
# ============================================================================
Model = DetectionModel  # retain YOLO 'Model' class for backwards compatibility